import argparse
import json
from pathlib import Path

def to_code(v):
    if isinstance(v, (int, float)):
        return str(int(v)).zfill(6)
    s = str(v).strip()
    digits = "".join(ch for ch in s if ch.isdigit())
    return (digits[:6] if digits else s[:6]).zfill(6)

def run(xlsx_path: str, json_path: str, sheet_index: int, code_col_index: int, name_col_index: int, header_rows: int):
    import openpyxl
    wb = openpyxl.load_workbook(xlsx_path, data_only=True)
    ws = wb.worksheets[sheet_index]
    rows = list(ws.iter_rows(values_only=True))
    mapping = {}
    start_idx = header_rows
    for row in rows[start_idx:]:
        if not row:
            continue
        code_val = row[code_col_index] if code_col_index < len(row) else None
        name_val = row[name_col_index] if name_col_index < len(row) else None
        if code_val is None or name_val is None:
            continue
        code = to_code(code_val)
        name = str(name_val).strip()
        if code and name:
            mapping[code] = name
    Path(json_path).write_text(json.dumps(mapping, ensure_ascii=False), encoding="utf-8")
    print(json.dumps({"status": "ok", "count": len(mapping), "out": json_path}, ensure_ascii=False))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("xlsx", nargs="?", default="app/static/A股列表.xlsx")
    parser.add_argument("json", nargs="?", default="app/static/stocks_cn.json")
    parser.add_argument("sheet", nargs="?", type=int, default=0)
    parser.add_argument("code_col", nargs="?", type=int, default=4)
    parser.add_argument("name_col", nargs="?", type=int, default=5)
    parser.add_argument("header_rows", nargs="?", type=int, default=1)
    args = parser.parse_args()
    run(args.xlsx, args.json, args.sheet, args.code_col, args.name_col, args.header_rows)

if __name__ == "__main__":
    main()